#!/usr/bin/env python
import flowd
import string
import getopt
import sys

mask_fields = [
	'TAG'			  ,
	'UNIX_SECS'       ,
    'UNIX_NSECS'      ,
    'SYSUPTIME'       ,
    'EXADDR'          ,   
    'DPKTS'           ,
    'DOCTETS'         ,
    'FIRST'           ,  
    'LAST'            ,
    'ENGINE_TYPE'     ,
    'ENGINE_ID'       ,    
    'SRCADDR'         ,
    'DSTADDR'         ,
    'NEXTHOP'         ,
    'INPUT'           ,
    'OUTPUT'          ,
    'SRCPORT'         ,   
    'DSTPORT'         ,
    'PROT'            ,
    'TOS'             ,
    'TCP_FLAGS'       ,    
    'SRC_MASK'        ,
    'DST_MASK'        ,
    'SRC_AS'          ,
    'DST_AS'          ,   
]

def format0(src,log):
	infile = open(src, "r")	
	flog = flowd.FlowLog(log, "wb")	
	mask = infile.readline().strip().split(',')
	for m in mask:
		if m.startswith('#:'):
			m = m.replace('#:','')
		if m.upper() not in mask_fields:
			print m.upper()
			print >> sys.stderr, "Unrecognized field:"+m
	lines = infile.readlines()
	count = 0
	for str in lines:
		data = str.strip().split(',')
		if len(mask) != len(data):
			print >> sys.stderr, "Invalid count of data"
		i=0
		flow = flowd.Flow()
		for flowfield in data:
			if mask[i] in ('TAG', 'tag'):
				flow.tag = int(flowfield)
			if mask[i] in ('UNIX_SECS', 'unix_secs'):
				flow.recv_sec = int(flowfield)
			if mask[i] in ('UNIX_NSECS', 'unix_nsecs'):
				flow.recv_usec = int(flowfield)
			if mask[i] in ('SYSUPTIME', 'sysuptime'):
				flow.sys_uptime_ms = int(flowfield)
			if mask[i] in ('EXADDR', 'exaddr'):
				flow.agent_addr = flowfield
			if mask[i] in ('DPKTS', 'dpkts'):
				flow.packets = int(flowfield)
			if mask[i] in ('DOCTETS', 'doctets'):
				flow.octets = int(flowfield)
			if mask[i] in ('FIRST', 'first'):
				flow.flow_start = int(flowfield)
			if mask[i] in ('LAST', 'last'):
				flow.flow_finish = int(flowfield)
			if mask[i] in ('ENGINE_TYPE', 'engine_type'):
				flow.engine_type = int(flowfield)
			if mask[i] in ('ENGINE_ID', 'engine_id'):
				flow.engine_id = int(flowfield)
			if mask[i] in ('SRCADDR', 'srcaddr'):
				flow.src_addr = flowfield
			if mask[i] in ('DSTADDR', 'dstaddr'):
				flow.dst_addr = flowfield
			if mask[i] in ('NEXTHOP', 'nexthop'):
				flow.gateway_addr = flowfield
			if mask[i] in ('INPUT', 'input'):
				flow.if_ndx_in = int(flowfield)
			if mask[i] in ('OUTPUT', 'output'):
				flow.if_ndx_out = int(flowfield)
			if mask[i] in ('SRCPORT', 'srcport'):
				flow.src_port = int(flowfield)
			if mask[i] in ('DSTPORT', 'dst_port'):
				flow.dst_port = int(flowfield)
			if mask[i] in ('PROT', 'prot'):
				flow.protocol = int(flowfield)
			if mask[i] in ('TOS', 'tos'):
				flow.tos = int(flowfield)
			if mask[i] in ('TCP_FLAGS', 'tcp_flags'):
				flow.tcp_flags = int(flowfield)
			if mask[i] in ('SRC_MASK', 'src_mask'):
				flow.src_mask = int(flowfield)
			if mask[i] in ('DST_MASK', 'dst_mask'):
				flow.dst_mask = int(flowfield)
			if mask[i] in ('SRC_AS', 'src_as'):
				flow.src_as = int(flowfield)
			if mask[i] in ('DST_AS', 'dst_as'):
				flow.dst_as = int(flowfield)
			flow.fields = 1074164750
			i += 1
		flog.write_flow(flow)
		count += 1
	print "imported %d flows"%count
	infile.close()

def usage():
	print >> sys.stderr, "import.py (flowd.py version %s)" % \
	    flowd.__version__
	print >> sys.stderr, "Usage: import.py [options] [source file] [flowd-store]";
	print >> sys.stderr, "Options:";
	print >> sys.stderr, "      -h       Display this help";
	print >> sys.stderr, "      -f       Format of output. 0 - ASCII";
	print >> sys.stderr, "      -m       mask_fields. Select fields from: %s. Default all fields" % ', '.join(mask_fields);
	sys.exit(1);
	
def main():

	try:
		opts, args = getopt.getopt(sys.argv[1:], 'f:m:u:h')
	except getopt.GetoptError:
		print >> sys.stderr, "Invalid commandline arguments"
		usage()

	format = None
	conn = None
	mask = None
	outfile = None
	uopt = None
	for o, a in opts:
		if o in ('-h', '--help'):
			usage()
			sys.exit(0)
		if o in ('-f', '--format'):
			format = a
		if o in('-m'):
			mask = a.split(',')
			for m in mask:
				if m not in mask_fields:
					print >> sys.stderr, "Unrecognized field:"+m
					sys.exit(1);
	if format is None:
		format = 0
	elif not format.isdigit() or int(format) > 0:
		print >> sys.stderr, "Invalid format"
		usage()
	if len(args) == 0:
		print >> sys.stderr, "No logfile specified"
		usage()
	elif len(args) != 2:
		print >> sys.stderr, "Arguments must consist of ascii source file and logfile to store data"
		usage()

	format0(args[0],args[1])	
			
if __name__ == '__main__': main()